import sys
import os.path
import pathlib

hs_ref_path = "/workspace/liuwei/public/reference/hg38/chip-seq/bwa_index/GRCh38_no_alt_analysis_set_GCA_000001405.15.fasta"
mm_ref_path = "/workspace/liuwei/public/reference/mm10/chip-seq/bwa_index/mm10_no_alt_analysis_set_ENCODE.fasta"
rn_ref_path = "/workspace/liuwei/public/reference/rn6/bwa-index/rn6.fasta"

HAS_CONTROL = True
control_bamfile_path = "control.sort.rmdup.bam"

def alignment(ref_path, reads_file1, reads_file2, sample_name):
    cmd_list = []
    cmd_list.append(f"bwa mem {ref_path} {reads_file1} {reads_file2} > {sample_name}.sam")
    return cmd_list

def process_samfile(sample_name):
    cmd_list = []
    cmd_list.append(f"samtools view -bS {sample_name}.sam -o {sample_name}.bam")
    cmd_list.append(f"samtools sort -o {sample_name}.sort.bam {sample_name}.bam")
    cmd_list.append(f"picard MarkDuplicates I={sample_name}.sort.bam O={sample_name}.sort.rmdup.bam M={sample_name}.sort.rmdup.metrics")
    cmd_list.append(f"samtools index {sample_name}.sort.rmdup.bam")
    return cmd_list

def call_peaks(species, sample_name):
    cmd_list = []
    if species in ["hs", "mm", "ce", "dm"]:
        gsize = species
    elif species == "rn":
        gsize = "2.75e+9"
    if HAS_CONTROL:
        cmd_list.append(call_peaks_use_macs2_with_control(gsize, sample_name, control_bamfile_path))
    else:
        cmd_list.append(call_peaks_use_macs2_without_control(gsize, sample_name))
    return cmd_list

def call_peaks_use_macs2_with_control(gsize, sample_name, control_name):
    cmd = f"macs2 callpeak -g {gsize} -p 1e-5 --keep-dup=auto -t {sample_name}.sort.rmdup.bam -c {control_name}.sort.rmdup.bam --outdir {sample_name}.macs.results -n {sample_name} &> {sample_name}.macs.log"
    return cmd

def call_peaks_use_macs2_without_control(gsize, sample_name):
    cmd = f"macs2 callpeak -g {gsize} -p 1e-5 --keep-dup=auto -t {sample_name}.sort.rmdup.bam --outdir {sample_name}.macs.results -n {sample_name} &> {sample_name}.macs.log"
    return cmd

def annotate_peaks(species, sample_name):
    cmd_list = []
    cmd_list.append(f"Rscript annotate_peak.r --peak_file {sample_name}.macs.results/{sample_name}_peaks.xls --anno_file {sample_name}.macs.results/{sample_name}.anno.xlsx --species {species}")
    return cmd_list


def build_pipeline(settings):
    cmd_list = []
    if settings['species'] == "mm":
        ref_path = mm_ref_path
    elif settings['species'] == "rn":
        ref_path = rn_ref_path
    elif settings['species'] == "hs":
        ref_path = hs_ref_path
    else:
        raise "species not found!"
    sample_name = settings['sample_name']
    species = settings['species']
    reads_file1 = settings['reads_file1']
    reads_file2 = settings['reads_file2']
    cmd_list += alignment(ref_path, reads_file1, reads_file2, sample_name)
    cmd_list += process_samfile(sample_name)
    cmd_list += call_peaks(species, sample_name)
    cmd_list += annotate_peaks(species, sample_name)
    pathlib.Path("scripts").mkdir(parents=True, exist_ok=True)
    script_file = os.path.join("scripts", f"script_{sample_name}.sh")
    with open(script_file, "w") as writer:
        for cmd in cmd_list:
            print(cmd, file=writer)
    return script_file

settings = [
    {'species': 'hs', 'reads_file1': 'sample_name.R1.fastq.gz', 'reads_file2': 'sample_name.R2.fastq.gz', 'sample_name': 'sample_name'},
]

if __name__ == "__main__":
    with open("pipeline.sh", "w") as writer:
        for sample_settings in settings:
            script_file = build_pipeline(sample_settings)
            print(f"bash {script_file}", file=writer)
